import time
import json
import random
import re
import websocket
from typing import Optional, Dict, Any, List, Union
from urllib.parse import urlparse
from locust import HttpUser, task, between, events

class WebSocketClient:
    def __init__(self):
        # 使用固定的WebSocket地址
        self.ws_url: str = "wss://game1.xy.siyuf.com/"
        self.ws: Optional[websocket.WebSocket] = None
        self.connected: bool = False
        self.user_id: Optional[str] = None
        self.token: Optional[str] = None
        self.character_id: Optional[str] = None
        self.server_id: Optional[str] = None
        self.server_list: List[Dict[str, Any]] = []
        self.game_state: str = "initial"  # 游戏状态：initial, server_selected, character_created, in_game
    
    def connect(self):
        """建立WebSocket连接"""
        try:
            print(f"尝试连接WebSocket: {self.ws_url}")
            self.ws = websocket.create_connection(self.ws_url)
            self.connected = True
            print("WebSocket连接成功")
            return True
        except Exception as e:
            print(f"WebSocket连接失败: {e}")
            return False
    
    def disconnect(self):
        """关闭WebSocket连接"""
        if self.ws and self.connected:
            self.ws.close()
            self.connected = False
            print("WebSocket连接已关闭")
    
    def send_message(self, cmd, data=None):
        """发送WebSocket消息"""
        if not self.connected or self.ws is None:
            print("WebSocket未连接，无法发送消息")
            return False
        
        message = {
            "cmd": cmd,
            "data": data if data is not None else {}
        }
        
        try:
            message_str = json.dumps(message)
            print(f"发送消息: {message_str}")
            self.ws.send(message_str)
            return True
        except Exception as e:
            print(f"发送消息失败: {e}")
            return False
    
    def receive_message(self, timeout=5):
        """接收WebSocket消息"""
        if not self.connected or self.ws is None:
            print("WebSocket未连接，无法接收消息")
            return None
        
        try:
            self.ws.settimeout(timeout)
            response = self.ws.recv()
            print(f"接收消息: {response[:200]}..." if len(response) > 200 else f"接收消息: {response}")
            return json.loads(response)
        except Exception as e:
            print(f"接收消息失败: {e}")
            return None
    
    def wait_for_specific_message(self, cmd_type, timeout=10, max_messages=10):
        """等待特定类型的消息"""
        start_time = time.time()
        messages_received = 0
        
        while time.time() - start_time < timeout and messages_received < max_messages:
            response = self.receive_message(timeout=2)
            messages_received += 1
            
            if response and response.get("cmd") == cmd_type:
                return response
        
        print(f"等待消息类型 {cmd_type} 超时")
        return None


class GameUser(HttpUser):
    wait_time = between(1, 3)  # 用户操作之间的等待时间
    host = "https://cdn.xy.siyuf.com"  # 游戏网站地址
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ws_client: WebSocketClient = WebSocketClient()  # 直接初始化
        self.username: str = ""
        self.password: str = ""
        self.html_content: Optional[str] = None
        self.server_list: List[Dict[str, Any]] = []
        self.selected_server: Optional[Dict[str, Any]] = None
    
    def on_start(self):
        """用户会话开始时执行"""
        self.ws_client = WebSocketClient()
        self.username = f"test_user_{random.randint(1000, 9999)}"
        self.password = "password123"
        self.html_content = None
        self.server_list = []
        self.selected_server = None
        
        # 步骤1: 模拟浏览器访问游戏网址
        self.visit_game_page()
        
        # 步骤2: 获取服务器列表
        self.get_server_list()
        
        # 步骤3-5: 连接WebSocket并完成游戏流程
        self.complete_game_flow()
    
    def on_stop(self):
        """用户会话结束时执行"""
        # 断开WebSocket连接
        if hasattr(self, 'ws_client') and self.ws_client and self.ws_client.connected:
            self.ws_client.disconnect()
    
    def visit_game_page(self):
        """步骤1: 模拟浏览器访问游戏网址"""
        try:
            print("步骤1: 访问游戏入口页面")
            with self.client.get("/XYres/index_waiwang.html?isPCWeiduan=true", catch_response=True) as response:
                if response.status_code == 200:
                    self.html_content = response.text
                    print("成功访问游戏入口页面")
                    
                    # 可以在这里提取一些关键信息，如CSRF令牌等
                    # 例如：self.csrf_token = re.search(r'csrf_token\s*=\s*["\']([^"\']+)["\']', self.html_content).group(1)
                else:
                    print(f"访问游戏入口页面失败: HTTP {response.status_code}")
        except Exception as e:
            print(f"访问游戏入口页面时发生错误: {e}")
    
    def get_server_list(self):
        """步骤2: 获取服务器列表"""
        try:
            print("步骤2: 获取服务器列表")
            # 这里的API路径需要根据实际情况调整
            with self.client.get("/api/server/list", catch_response=True) as response:
                if response.status_code == 200:
                    try:
                        result = response.json()
                        if "data" in result and isinstance(result["data"], list):
                            self.server_list = result["data"]
                            print(f"成功获取服务器列表，共 {len(self.server_list)} 个服务器")
                            
                            # 选择一个服务器（这里简单地选择第一个）
                            if self.server_list:
                                self.selected_server = self.server_list[0]
                                print(f"选择服务器: {self.selected_server.get('name', 'Unknown')}")
                        else:
                            print(f"服务器列表格式异常: {result}")
                    except Exception as e:
                        print(f"解析服务器列表失败: {e}")
                        print(f"响应内容: {response.text[:200]}...")
                else:
                    print(f"获取服务器列表失败: HTTP {response.status_code}")
        except Exception as e:
            print(f"获取服务器列表时发生错误: {e}")
            
            # 如果API调用失败，创建一个模拟的服务器列表用于测试
            self.server_list = [{"id": "1", "name": "测试服务器1", "status": "online"}]
            self.selected_server = self.server_list[0]
            print("使用模拟的服务器列表继续测试")
    
    def complete_game_flow(self):
        """步骤3-5: 连接WebSocket并完成游戏流程"""
        # 步骤3: 连接WebSocket
        if not self.ws_client.connect():
            print("WebSocket连接失败，无法继续游戏流程")
            return
        
        # 发送初始化消息
        self.ws_client.send_message("init", {
            "client_version": "1.0.0",
            "platform": "pc",
            "channel": "official"
        })
        
        # 等待初始化响应
        init_response = self.ws_client.wait_for_specific_message("init_response")
        if not init_response or not init_response.get("success", False):
            print("初始化失败，无法继续游戏流程")
            return
        
        # 步骤3: 选择服务器
        self.select_server()
        
        # 步骤4: 登录游戏
        if not self.login_game():
            return
        
        # 创建或选择角色
        if not self.select_or_create_character():
            return
        
        # 步骤5: 进入游戏地图
        self.enter_game_world()
    
    def select_server(self):
        """选择游戏服务器"""
        print("步骤3: 选择游戏服务器")
        if not self.selected_server:
            print("没有可用的服务器，无法继续")
            return False
        
        server_id = self.selected_server.get("id", "1")
        self.ws_client.send_message("select_server", {
            "server_id": server_id
        })
        
        # 等待服务器选择响应
        response = self.ws_client.wait_for_specific_message("server_selected")
        if response and response.get("success", False):
            # 使用字符串类型的server_id
            if server_id is not None:
                self.ws_client.server_id = str(server_id)
            self.ws_client.game_state = "server_selected"
            print(f"成功选择服务器: {self.selected_server.get('name', 'Unknown')}")
            return True
        else:
            print("选择服务器失败")
            return False
    
    def login_game(self):
        """登录游戏"""
        print("步骤4: 登录游戏")
        # 这里可以是游客登录或账号登录
        # 游客登录示例
        self.ws_client.send_message("login", {
            "login_type": "guest",
            "device_id": f"device_{random.randint(10000, 99999)}"
        })
        
        # 等待登录响应
        response = self.ws_client.wait_for_specific_message("login_response")
        if response and response.get("success", False):
            user_data = response.get("data", {})
            self.ws_client.user_id = user_data.get("user_id")
            self.ws_client.token = user_data.get("token")
            print(f"登录成功，用户ID: {self.ws_client.user_id}")
            return True
        else:
            print("登录失败")
            return False
    
    def select_or_create_character(self):
        """选择或创建角色"""
        print("获取角色列表")
        self.ws_client.send_message("get_characters", {})
        
        # 等待角色列表响应
        response = self.ws_client.wait_for_specific_message("character_list")
        if not response:
            print("获取角色列表失败")
            return False
        
        characters = response.get("data", {}).get("characters", [])
        
        if characters:
            # 选择现有角色
            selected_character = characters[0]
            self.ws_client.character_id = selected_character.get("id")
            print(f"选择现有角色: {selected_character.get('name')}")
        else:
            # 创建新角色
            print("创建新角色")
            character_name = f"hero_{random.randint(1000, 9999)}"
            character_class = random.choice(["warrior", "mage", "archer", "priest"])
            
            self.ws_client.send_message("create_character", {
                "name": character_name,
                "class": character_class,
                "gender": random.choice(["male", "female"]),
                "appearance": {
                    "hair": random.randint(1, 10),
                    "face": random.randint(1, 10),
                    "color": random.randint(1, 10)
                }
            })
            
            # 等待角色创建响应
            response = self.ws_client.wait_for_specific_message("character_created")
            if response and response.get("success", False):
                self.ws_client.character_id = response.get("data", {}).get("character_id")
                self.ws_client.game_state = "character_created"
                print(f"角色创建成功: {character_name}, ID: {self.ws_client.character_id}")
            else:
                print("角色创建失败")
                return False
        
        return True
    
    def enter_game_world(self):
        """步骤5: 进入游戏地图"""
        print("步骤5: 进入游戏地图")
        if not self.ws_client.character_id:
            print("没有角色ID，无法进入游戏")
            return False
        
        self.ws_client.send_message("enter_game", {
            "character_id": self.ws_client.character_id
        })
        
        # 等待进入游戏响应
        response = self.ws_client.wait_for_specific_message("game_entered")
        if response and response.get("success", False):
            self.ws_client.game_state = "in_game"
            print("成功进入游戏世界")
            
            # 处理初始化数据
            game_data = response.get("data", {})
            position = game_data.get("position", {})
            print(f"角色位置: X={position.get('x', 0)}, Y={position.get('y', 0)}, 地图ID={game_data.get('map_id', 'unknown')}")
            
            # 可以在这里处理更多的初始化数据
            # 例如：背包物品、技能列表、任务列表等
            
            return True
        else:
            print("进入游戏世界失败")
            return False
    
    @task
    def game_actions(self):
        """游戏中的各种操作"""
        # 只有在游戏中才执行操作
        if self.ws_client.game_state != "in_game":
            return
        
        # 随机选择一个操作执行
        action = random.choice(["move", "chat", "check_surroundings"])
        
        if action == "move":
            self.move_character()
        elif action == "chat":
            self.send_chat_message()
        elif action == "check_surroundings":
            self.check_surroundings()
    
    def move_character(self):
        """移动角色"""
        # 随机移动方向和距离
        direction = random.choice(["north", "south", "east", "west"])
        distance = random.randint(1, 10)
        
        self.ws_client.send_message("move", {
            "direction": direction,
            "distance": distance
        })
        
        response = self.ws_client.wait_for_specific_message("move_result")
        if response and response.get("success", False):
            new_position = response.get("data", {}).get("position", {})
            print(f"角色移动成功，新位置: X={new_position.get('x')}, Y={new_position.get('y')}")
        else:
            print("角色移动失败")
    
    def send_chat_message(self):
        """发送聊天消息"""
        messages = [
            "你好！", "有人组队吗？", "这个游戏真好玩！", 
            "请问新手任务在哪里？", "有人知道怎么升级快吗？"
        ]
        
        self.ws_client.send_message("chat", {
            "channel": "world",
            "message": random.choice(messages)
        })
        
        # 聊天消息通常不需要等待响应
        print("聊天消息已发送")
    
    def check_surroundings(self):
        """查看周围环境"""
        self.ws_client.send_message("get_surroundings", {
            "range": random.randint(5, 20)
        })
        
        response = self.ws_client.wait_for_specific_message("surroundings")
        if response and response.get("success", False):
            entities = response.get("data", {}).get("entities", [])
            print(f"周围有 {len(entities)} 个实体")
        else:
            print("获取周围环境失败")


# 自定义事件处理
@events.init.add_listener
def on_locust_init(environment, **kwargs):
    print("Locust初始化完成，准备开始测试...")

@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    print("测试开始...")

@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    print("测试结束...")


if __name__ == "__main__":
    # 本地测试用
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "debug":
        # 创建一个测试用户并执行完整流程
        # 使用环境对象代替None
        from locust.env import Environment
        env = Environment(user_classes=[GameUser])
        user = GameUser(env)
        
        user.on_start()
        
        # 如果成功进入游戏，执行一些游戏操作
        if user.ws_client and user.ws_client.game_state == "in_game":
            for _ in range(5):
                user.move_character()
                time.sleep(1)
                user.send_chat_message()
                time.sleep(1)
                user.check_surroundings()
                time.sleep(1)
        
        user.on_stop()